# @Author：zh
# @Data：2021/12/28 21:14
# @：PyCharm
# Python版本：3.7

## 基于感知机模型预测
import numpy as np
import torch
from torch import nn
import  cv2 as cv


def data_tf(x):
    x = np.array(x, dtype='float32') / 255
    x = (x - 0.5) / 0.5 # 标准化，这个技巧之后会讲到
    x = x.reshape((-1,)) # 拉平成一个一维向量
    ##  numpy转换为tensor
    x = torch.from_numpy(x)
    return x

net = nn.Sequential(
    ## 全连接层：输入 输出
    nn.Linear(784, 400),
    nn.ReLU(),
    nn.Linear(400, 200),
    nn.ReLU(),
    nn.Linear(200, 100),
    nn.ReLU(),
    nn.Linear(100, 10),##最后输出10个类别
    # nn.Softmax(dim=1)
)

img=cv.imread("data/TestData/5.jpg",)
img=cv.cvtColor(img, cv.COLOR_BGR2GRAY)
img=cv.resize(img,(28,28))
array=np.array(img)
tensor=data_tf(array)

tensor = tensor.view(1,784)

model = net(tensor.to(torch.float32))
net.load_state_dict(torch.load("Model_Minst_SoftMax.pt"))
net.eval()

output = net(tensor.to(torch.float32))
_, prediction = torch.max(output, 1)
#将预测结果从tensor转为array，并抽取结果
prediction = prediction.numpy()[0]

print(prediction)
